{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a RAG Pipeline over Legal Documents\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_parse/blob/main/examples/multimodal/legal_rag.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "This example shows how LlamaParse and LlamaIndex can be used to parse various types of legal documents, which may contain complex tabular data. The advantage of this is being able to quickly retrieve a specific answer to a legal question with comprehensive context — knowledge of precedents, statutes, and cases presented in the given documents. A user can quickly find the answer to or find out more details about a specific legal question without having to read through the often long documents by using LLMs.\n",
    "\n",
    "In this example, we will be using legal documents from the archive of the Library of Congress ([link to dataset](https://www.loc.gov/item/2020445568/)). These documents vary by format, with some containing pure text and others containing headings, sections, and large tables. This shows how LlamaParse can parse a wide variety of documents and still retrieve accurate results.\n",
    "\n",
    "The documents in this example include:\n",
    "- [APA Program Report](https://www.irs.gov/pub/irs-apa/a_2003-19.pdf)\n",
    "- [2004 Report on the CRA performance of Barre Savings Bank in Barre, MA](https://github.com/user-attachments/files/16536412/barre_savings_bank_evaluation.pdf)\n",
    "- [2016 Energy Supply/Demand Forecast](https://github.com/user-attachments/files/16536415/energy_supply_demand.pdf)\n",
    "- [Transcript of Senate Committee Hearing about Foreign Markets](https://github.com/user-attachments/files/16536422/foreign_markets.pdf)\n",
    "- [A Motion To Stay for an Indiana Court Case](https://github.com/user-attachments/files/16536427/motion_to_stay.pdf)\n",
    "- [Article About an OC Representative's Bill to Introduce Offshore Drilling to CA](https://github.com/user-attachments/files/16536437/oc_bill_offshore_drilling.pdf)\n",
    "- [Charter of the Subcommittee on Ocean Science and Technology](https://github.com/user-attachments/files/16536445/ost_subcommittee_charter.pdf)\n",
    "- [US Immigration Case](https://github.com/user-attachments/files/16536446/us_immigration_case.pdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup and Installation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install LlamaIndex, download the data, and apply `nest_asyncio`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index llama-parse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://github.com/user-attachments/files/16447759/data.zip -O data.zip\n",
    "!unzip -o data.zip\n",
    "!rm data.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up your OpenAI and LlamaCloud keys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"<Your OpenAI API Key>\"\n",
    "os.environ[\"LLAMA_CLOUD_API_KEY\"] = \"<Your LlamaCloud API Key>\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code Implementation\n",
    "\n",
    "Set up LlamaParse. We want to parse the PDF files into markdown, translating the tabular data into markdown tables. To ensure accuracy, we will use the GPT-4o multimodal model to parse the PDFs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_parse import LlamaParse\n",
    "\n",
    "parser = LlamaParse(\n",
    "    result_type=\"markdown\",\n",
    "    parsing_instruction=\"Provided are a series of US legal documents.\",\n",
    "    use_vendor_multimodal_model=True,\n",
    "    vendor_multimodal_model_name=\"openai-gpt4o\",\n",
    "    show_progress=True,\n",
    ")\n",
    "\n",
    "DATA_DIR = \"data\"\n",
    "\n",
    "\n",
    "def get_data_files(data_dir=DATA_DIR) -> list[str]:\n",
    "    files = []\n",
    "    for f in os.listdir(data_dir):\n",
    "        fname = os.path.join(data_dir, f)\n",
    "        if os.path.isfile(fname):\n",
    "            files.append(fname)\n",
    "    return files\n",
    "\n",
    "\n",
    "files = get_data_files()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load data from parser into documents containing parsed Markdown text from the legal document PDFs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Parsing files: 100%|██████████| 8/8 [01:25<00:00, 10.67s/it]\n"
     ]
    }
   ],
   "source": [
    "documents = parser.load_data(\n",
    "    files,\n",
    "    extra_info={\"name\": \"US legal documents provided by the Library of Congress.\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup LlamaIndex. Set the default LLM to GPT-4o (a multi-modal model), and create an index from the documents, and persist these documents to disk. If these documents have already been persisted, then load index from the persisted docs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import (\n",
    "    VectorStoreIndex,\n",
    "    StorageContext,\n",
    "    load_index_from_storage,\n",
    "    Settings,\n",
    ")\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "embed_model = OpenAIEmbedding(model=\"text-embedding-3-large\")\n",
    "llm = OpenAI(\"gpt-4o\")\n",
    "\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model\n",
    "\n",
    "if not os.path.exists(\"storage_legal\"):\n",
    "    index = VectorStoreIndex(documents, embed_model=embed_model)\n",
    "    index.storage_context.persist(persist_dir=\"./storage_legal\")\n",
    "else:\n",
    "    ctx = StorageContext.from_defaults(persist_dir=\"./storage_legal\")\n",
    "    index = load_index_from_storage(ctx)\n",
    "\n",
    "query_engine = index.as_query_engine()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example Queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "The majority of Barre Savings Bank's loans went to residential real estate, specifically 1-4 family mortgages, which accounted for 78.7 percent of the total loans."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import display, Markdown\n",
    "\n",
    "response = query_engine.query(\n",
    "    \"Where did the majority of Barre Savings Bank's loans go?\"\n",
    ")\n",
    "display(Markdown(str(response)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "Mr. Kubarych believes foreign markets are important because they are attractive to foreign investors for the same reasons they are attractive to Americans. The economic data is strong, and the high tech boom has created a positive perception that overshadows longer-term vulnerabilities. Additionally, foreign investors have high expectations for the U.S. to maintain a firm monetary policy in response to inflation and to act as a superpower rather than pursuing narrow nationalist economic policies."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = query_engine.query(\n",
    "    \"Why does Mr. Kubarych believe foreign markets are so important?\"\n",
    ")\n",
    "display(Markdown(str(response)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "House Speaker Nancy Pelosi and the Democratic majority are against the proposal of offshore drilling in California. Pelosi stated that offshore drilling is \"off the table,\" and Democrats have been consistently unwilling to bend environmental rules. They argue that oil companies are not using the 68 million acres of federal lands already leased to them, either because it takes a long time or they lack the necessary equipment."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = query_engine.query(\n",
    "    \"Who is against the proposal of offshore drilling in CA and why?\"\n",
    ")\n",
    "display(Markdown(str(response)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "The purpose of the Ocean Science and Technology Subcommittee (SOST) is to advise and assist the Committee on Environment, Natural Resources, and Sustainability on national issues of ocean science and technology. The SOST aims to contribute to the goals for Federal ocean science and technology by developing coordinated interagency strategies. It also retains the functions of the previously-chartered Joint Subcommittee on Ocean Science and Technology and serves as the Ocean Science and Technology Interagency Policy Committee for the National Ocean Council."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = query_engine.query(\n",
    "    \"What is the purpose of the Ocean Science and Technology Subcommittee?\"\n",
    ")\n",
    "display(Markdown(str(response)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "The immigration appeal is dismissed because the petitioner is not a U.S. citizen, and therefore, is not eligible to file a Petition for Alien Fiancé(e) (Form I-129F) on behalf of the beneficiary. The relevant law provides nonimmigrant classification only to aliens who are the fiancé(e)s of U.S. citizens."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = query_engine.query(\"Why is the immigration appeal dismissed?\")\n",
    "display(Markdown(str(response)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "An advance pricing agreement (APA) is a binding contract between a taxpayer and the IRS that establishes an approved transfer pricing method (TPM) for specific transactions. This agreement aims to prevent disputes over transfer pricing by ensuring that the taxpayer's tax returns for the covered years are consistent with the agreed TPM. APAs can be unilateral, involving only the taxpayer and the IRS, or bilateral/multilateral, involving agreements with one or more foreign tax authorities to avoid double taxation."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = query_engine.query(\"What is an advance pricing agreement?\")\n",
    "display(Markdown(str(response)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-parse-5ZmnAQ0r-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
